//
// Created by myrica on 3/30/25.
//

#include "Pass/RemoveUnreachableOperation.h"
#include <ranges>

std::unique_ptr<mlir::Pass> irispass::createRemoveUnreachableOperationPass() {
    return std::make_unique<RemoveUnreachableOperation>();
}

void irispass::RemoveUnreachableOperation::runOnOperation() {
    mlir::ModuleOp module = getOperation();

    module.walk([&](mlir::func::FuncOp func) {
        llvm::errs() << "Current function: " << func.getName() << "\n";

        if (func.getBody().empty())
            return;

        int deleted = 0;
        func.walk([&](mlir::Block *block) {
            if (block->getParentOp() != func)
                return;

            auto term_iter = std::ranges::find_if(block->getOperations(),
                                                  [&](mlir::Operation &op) {
                                                      return op.hasTrait<mlir::OpTrait::IsTerminator>();
                                                  });
            assert(term_iter != block->getOperations().end());

            auto _erase_self_and_uses_recursively = [&](this auto &&self, mlir::Operation &iter) -> void {
                if (iter.getUses().empty()) {
                    iter.erase();
                    deleted++;
                } else {
                    for (auto &use: iter.getUses()) {
                        self(*use.getOwner());
                    }
                }
            };

            do {
                auto next_node = term_iter->getNextNode();
                if (next_node == nullptr)
                    break;
                auto invalid = next_node->getIterator();
                _erase_self_and_uses_recursively(*invalid);
            } while (term_iter->getNextNode() != nullptr && term_iter->getNextNode()->getIterator() != block->
                     getOperations().end());
        });
        llvm::errs() << "Removed " << deleted << " unreachable operation(s) in function " << func.getName() << "\n";
    });
}
